from keras.models import Model
from keras.utils import plot_model
import numpy as np
import innvestigate.utils as iutils
import os
from keras.preprocessing.image import img_to_array, load_img
import tensorflow as tf
import PIL.Image
import matplotlib.pyplot as plt
import cv2
import pdb
from interpreters import GradCAM, SGLRP
from utils import *
import json
  
plt.switch_backend('agg')

def visualize_gradcam(analysis, src=None, size=224):
    if np.max(src) > 1:
        raise Exception(
            "The input image should np.float32 in the range [0, 1]")
    colormap = cv2.COLORMAP_JET
    mask = (analysis - analysis.min()) / analysis.max()
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)[:, :, ::-1]
    heatmap = np.float32(heatmap) / 255

    cam = heatmap + src
    cam = cam / np.max(cam)
    return np.uint8(255 * cam)


def build_gradcam(partial_model, out=None, out_ori=None, layer_name=None):
    assert layer_name is not None
    if out is None: 
        assert out_ori is None
        out = partial_model.output
    if out_ori is not None: 
        indics_ori = tf.argmax(out_ori[0])
        out_modified = tf.where(tf.equal(tf.range(1000), tf.cast(indics_ori, tf.int32)), tf.zeros(1000), out[0])
        target_id = tf.argmax(out_modified)
    else:
        target_id = tf.argmax(out[0])
    
    analyzer = GradCAM(partial_model, target_id=target_id, relu=True, layer_name=layer_name)    
    return analyzer


def test_gradcam():
    os.environ["CUDA_VISIBLE_DEVICES"] = '2'
    layers_dict = parse_json()
    for net_name in ['VGG19', 'ResNet50', 'DenseNet201']:
        model, preprocess_input = load_net(net_name) 
        
        dot_img_file = "demo/Gradcam/{}_model_arch.png".format(net_name)
        plot_model(model, to_file=dot_img_file, show_shapes=True)
        analyser = build_gradcam(Model(inputs=model.inputs, outputs=iutils.keras.graph.pre_softmax_tensors(model.outputs), name=model.name + '_partial'), layer_name=layers_dict[net_name])
        
        image_size = int(model.input.shape[1])
        image_imagenet_ = process_sample('demo/ImageNet_ILSVRC2012_val_00000046.JPEG', image_size)
        image_damagenet_ = process_sample('demo/DAmageNet_ILSVRC2012_val_00000046.png', image_size)
        image_imagenet = preprocess_input(image_imagenet_.copy())
        image_damagenet = preprocess_input(image_damagenet_.copy())
        
        plot = Plot('.', n_img_x=2, img_w=224, img_h=224, img_c=3)
        clean_att = analyser.analyze([image_imagenet])
        adv_att = analyser.analyze([image_damagenet])
        
        plot.add_image(visualize_gradcam(clean_att, src=image_imagenet_ / 255))
        plot.add_image(visualize_gradcam(adv_att, src=image_damagenet_ / 255))
        img_path = 'demo/Gradcam/' + net_name + '.png'
        plot.save_images(img_path)
        print('\n', img_path, '\n')


if __name__ == '__main__':
    test_gradcam()